import os
import numpy as np

def get_wafer_map(wafer_text, wafer_folder, predict_data):

    WaferMapRecipe = os.path.join(wafer_folder, "WaferMapRecipe.ini")
    ColorImageGrabingInfo = os.path.join(wafer_folder, "ColorImageGrabingInfo.ini")
    ProductInfo = os.path.join(wafer_folder, "ProductInfo.ini")

    with open(WaferMapRecipe) as f:
        lines = f.readlines()

        for line in lines:
            if "CustPoint2.col" in line:
                cp_col = int(line.strip().split("=")[-1])
            if "CustPoint2.row" in line:
                cp_row = int(line.strip().split("=")[-1])
            if "RefPoint1.col" in line:
                rp_col = int(line.strip().split("=")[-1])
            if "RefPoint1.row" in line:
                rp_row = int(line.strip().split("=")[-1])
    with open(ProductInfo) as f:
            lines = f.readlines()

            for line in lines:
                if "XDieIndex" in line:
                    die_x = float(line.strip().split("=")[-1])
                if "YDieIndex" in line:
                    die_y = float(line.strip().split("=")[-1])

    # 读取 wafer txt 文件
    with open(wafer_text) as f:
        lines = f.readlines()

    # 从 txt 文件中读取 waferMapDim
    row_num = int(lines[4].strip()[6:])
    col_num = int(lines[5].strip()[6:])

    lines = lines[12:]

    wafer_data = np.zeros((row_num, col_num), dtype=int).tolist()

    for i in range(len(lines)):
        line = lines[i]
        line = line[8:]
        line_strings = line.split()

        for j in range(len(line_strings)):
            line_string = line_strings[j]

            if line_string=="___":
                line_string = 0
            else:
                line_string = 1
            wafer_data[i][j] = line_string


    # 每一张图片的推理结果
    for image_predict in predict_data:

        filename = image_predict["filename"]
        col = int(image_predict["Col"])
        row = int(image_predict["Row"])

        wuran_box = image_predict["wuran_box"]
        wuran_crop = image_predict.get("wuran_crop", [])
        zhenhen_box = image_predict["zhenhen_box"]
        zhenhen_crop = image_predict.get("zhenhen_crop", [])
        # zhenhen_num_error = image_predict.get("zhenhen_num_error", [])

        can_wuran_box = image_predict.get("can_wuran_box", [])
        madian_box = image_predict.get("madian_box", [])

        # 真正绘制 wafer 的行和列 
        txt_col = col-(rp_col-cp_col)-1
        txt_row = row-(rp_row-cp_row)-1

        if wafer_data[txt_row][txt_col] == 1:
            wafer_data[txt_row][txt_col] = []

        for box in wuran_box:
            fault_type = 2
            fault_x = box[1][0]
            fault_y = box[1][1]

            # 计算错误点在一个 die 上的相对位置
            relative_x = fault_x - die_x*col
            relative_y = fault_y - die_y*row

            wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})

        for crop in wuran_crop:
            fault_type = 2
            fault_x = crop[1][0]
            fault_y = crop[1][1]

            # 计算错误点在一个 die 上的相对位置
            relative_x = fault_x - die_x*col
            relative_y = fault_y - die_y*row

            wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})

        for box in zhenhen_box:
            fault_type = 3
            fault_x = box[1][0]
            fault_y = box[1][1]

            # 计算错误点在一个 die 上的相对位置
            relative_x = fault_x - die_x*col
            relative_y = fault_y - die_y*row

            wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})

        for crop in zhenhen_crop:
            fault_type = 3
            fault_x = crop[1][0]
            fault_y = crop[1][1]

            # 计算错误点在一个 die 上的相对位置
            relative_x = fault_x - die_x*col
            relative_y = fault_y - die_y*row

            wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})

        # for error in zhenhen_num_error:
        #     fault_type = 4
        #     fault_x = error[1][0]
        #     fault_y = error[1][1]

        #     # 计算错误点在一个 die 上的相对位置
        #     relative_x = fault_x - die_x*col
        #     relative_y = fault_y - die_y*row

        #     wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})
        for box in can_wuran_box:
            fault_type = 5
            fault_x = box[1][0]
            fault_y = box[1][1]

            # 计算错误点在一个 die 上的相对位置
            relative_x = fault_x - die_x*col
            relative_y = fault_y - die_y*row

            wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})

        for box in madian_box:
            fault_type = 6
            fault_x = box[1][0]
            fault_y = box[1][1]

            # 计算错误点在一个 die 上的相对位置
            relative_x = fault_x - die_x*col
            relative_y = fault_y - die_y*row

            wafer_data[txt_row][txt_col].append({"fault_type": fault_type, "relative_fault_x": relative_x/die_x, "relative_fault_y": relative_y/die_y, "image_path": os.path.join(wafer_folder, filename), "fault_x": fault_x, "fault_y": fault_y})

        # 如果模型推理是个正确图片
        if wafer_data[txt_row][txt_col] == []:
            wafer_data[txt_row][txt_col] = 1

    return wafer_data